/**************************************************************************
* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
*                                                                        *
* Author: The ALICE Off-line Project.                                    *
* Contributors are mentioned in the code where appropriate.              *
*                                                                        *
* Permission to use, copy, modify and distribute this software and its   *
* documentation strictly for non-commercial purposes is hereby granted   *
* without fee, provided that the above copyright notice appears in all   *
* copies and that both the copyright notice and this permission notice   *
* appear in the supporting documentation. The authors make no claims     *
* about the suitability of this software for any purpose. It is          *
* provided "as is" without express or implied warranty.                  *
**************************************************************************/
//
//  Xian-Guo Lu <lu@physi.uni-heidelberg.de>
//
#include "TH1I.h"
#include "TH2D.h"
#include "TParticle.h"
#include "TAxis.h"

#include "AliAnalysisManager.h"
#include "AliCentrality.h"
#include "AliESDInputHandler.h"
#include "AliPIDResponse.h"
#include "AliESDtrackCuts.h"
#include "AliESDv0.h"
#include "AliESDv0KineCuts.h"
#include "AliMCEventHandler.h"
#include "AliMCEvent.h"
#include "AliTPCdEdxInfo.h"
#include "AliStack.h"

#include "AliAnalysisTaskpkRaa.h"

const Double_t EPSILON = 1e-12;

Int_t AliAnalysisTaskpkRaa::SetCenAxis(TAxis *ax)
{
  Double_t cens[]={0, 5, 10, 20, 40, 60, 80, 100};
  const Int_t ncen = sizeof(cens)/sizeof(Double_t)-1;

  ax->Set(ncen, cens);
  return ncen;
}

AliAnalysisTaskpkRaa::AliAnalysisTaskpkRaa():
  AliAnalysisTaskSE("pkRaa")
  , fMCEvent(0x0)
  , fMCStack(0x0)
  , fMCStream(0x0)
  , fMCTree(0x0)
  , fEvent(0x0)
  , fpid(0x0)
  , fStream(0x0)
  , fTree(0x0)
  , fList(0x0)
  , fCuts(0x0)
  , fVtxZ(0x0)
  , fNsel(0x0)
  , fCen(0x0)
  , fEta(0x0)
  , fNevent(0x0)
{
  //
  //
  //

  DefineOutput(1, TTree::Class());
  DefineOutput(2, TList::Class());
  DefineOutput(3, TTree::Class());
  
  fMCStream = new TTreeStream("mctree");
  fMCTree=fMCStream->GetTree();

  fStream = new TTreeStream("tree");
  fTree=fStream->GetTree();

  fList = new TList;
  fList->SetOwner(kTRUE);

  fCuts = new TH2D("cuts","", 31,-0.5,30.5, 7,0,100);
  SetCenAxis(fCuts->GetYaxis());

  fVtxZ = new TH1D("vtxZ","",500, -50,50);
  fNsel = new TH1I("nsel","",100,0,100);
  fCen = new TH1D("cen","",110,-10,100);
  fEta = new TH1D("eta","",100,-1,1);
  fNevent = new TH1I("nevent","",2,0,2);

  fList->Add(fCuts);
  fList->Add(fVtxZ);
  fList->Add(fNsel);
  fList->Add(fCen);
  fList->Add(fEta);
  fList->Add(fNevent);
}

AliAnalysisTaskpkRaa::~AliAnalysisTaskpkRaa()
{
  //
  // Destructor
  //

  /* can not be deleted in chain, they are in input handler!
  delete fEvent;
  */

  //delete fStream;
  delete fList;
}

void AliAnalysisTaskpkRaa::UserCreateOutputObjects()
{
  //
  //
  //

  PostData(1, fTree);
  PostData(2, fList);
  PostData(3, fMCTree);
}

void AliAnalysisTaskpkRaa::UserExec(Option_t *)
{
  //
  //Execute
  // 
  Float_t v0cen = -999;

  fCuts->Fill(0., v0cen); //================================

  if(!CutInput())
    return;

  fCuts->Fill(1., v0cen); //================================

  const AliCentrality *evtCentrality = fEvent->GetCentrality();
  if(evtCentrality){
    v0cen = evtCentrality->GetCentralityPercentile("V0M");
  }
  fCen->Fill(v0cen);

  fCuts->Fill(2., v0cen); //================================

  const Int_t kFailEvent = FailCutEvent(v0cen);

  //as is observed in output, cut 0 and 1 is equivalent
  if(fMCStack){
    SaveMCgen(v0cen, kFailEvent);
  }

  if(kFailEvent){
    return;
  }

  fCuts->Fill(3., v0cen); //================================

  const Int_t ntrk0 = fEvent->GetNumberOfTracks();
  typetrack * trks[ntrk0];
  for(Int_t ii=0; ii<ntrk0; ii++){
    trks[ii] = 0x0;
  }
  Int_t cutbit[ntrk0];
  Int_t nsel = 0;
  CutTrack(nsel, trks, cutbit);
  fNsel->Fill(nsel);
  if(!nsel)
    return;

  fCuts->Fill(4., v0cen); //================================

  Int_t ifilled = 0;
  for(Int_t itrk=0; itrk<nsel ; itrk++){
    typetrack * trackptr = trks[itrk];
   
    Float_t eta = trackptr->Eta();
    fEta->Fill(eta);
    if(!EtaCut(eta))
      continue;

    Float_t mom = trackptr->P();
    Int_t charge = trackptr->Charge();
    Float_t pt = trackptr->Pt();
   
    Float_t dcaxy = -999, dcaz = -999;
    trackptr->GetImpactParameters(dcaxy, dcaz);

    Float_t tpcchi2cg = trackptr->GetChi2TPCConstrainedVsGlobal(fEvent->GetPrimaryVertexTracks());

    Int_t tpcncross = trackptr->GetTPCCrossedRows();
    Float_t tpcsig  = trackptr->GetTPCsignal();
    Float_t tpcip = trackptr->GetTPCmomentum();
    Int_t tpcnf = trackptr->GetTPCNclsF();
    Int_t tpcpidn = trackptr->GetTPCsignalN();
    Int_t tpcns = trackptr->GetTPCnclsS();
    Int_t tpcncls = trackptr->GetTPCncls();

    //TOF
    //quantity during reconstruction not good to use
    //any way the resolution too bad and peak position not correct
    /*
    Float_t tofmass = -999;
    if(trackptr->GetStatus()&AliESDtrack::kTOFpid){
      if((trackptr->GetIntegratedLength() != 0) && (trackptr->GetTOFsignal() != 0)){
        const Double_t tofbeta = trackptr->GetIntegratedLength()/trackptr->GetTOFsignal()/2.99792458e-2;
        if(tofbeta<1 && tofbeta> EPSILON){
          const Double_t tofbg = tofbeta/TMath::Sqrt(1-TMath::Power(tofbeta,2));
          tofmass = mom/tofbg;
        }
      }
    }
    */

    //===>
    Float_t tofpion = -999;
    Float_t tofkaon = -999;
    Float_t tofproton = -999;
    Float_t tofelectron = -999;

    //http://alisoft.cern.ch/viewvc/trunk/STEER/STEERBase/AliPIDResponse.cxx?view=markup&root=AliRoot
    //check kTOFout, kTIME, mismatch
    const Int_t dummynSpe = 4;
    Double_t dummypro[4]={-999,-999,-999,-999};
    AliPIDResponse::EDetPidStatus tofstatus = fpid->ComputeTOFProbability(trackptr, dummynSpe, dummypro);
    if( tofstatus == AliPIDResponse::kDetPidOk ){
      tofpion     = fpid->NumberOfSigmasTOF( trackptr, AliPID::kPion);
      tofkaon     = fpid->NumberOfSigmasTOF( trackptr, AliPID::kKaon);
      tofproton   = fpid->NumberOfSigmasTOF( trackptr, AliPID::kProton);
      tofelectron = fpid->NumberOfSigmasTOF( trackptr, AliPID::kElectron);
    }
    //<===

    //-------------------------------------------------- MC
    Int_t mcpdg = -999;
    Int_t mctype = 0;
    Float_t mcpt = -999;
    Float_t mceta = -999;
    Int_t mccharge = -999;
    if(fMCStack){
      const Int_t mclabel = TMath::Abs(trackptr->GetLabel());
      TParticle *trackMC = fMCStack->Particle(mclabel);
      mcpdg = TMath::Abs(trackMC->GetPdgCode());
      mcpt = trackMC->Pt();
      mceta = trackMC->Eta();

      //root [11] TParticlePDG::Charge(
      //Double_t Charge() const         // charge in units of |e|/3
      if(trackMC->GetPDG()->Charge()>0)
        mccharge = 1;
      else if(trackMC->GetPDG()->Charge()<0)
        mccharge = -1;
      else 
        mccharge = 0;

      if(fMCStack->IsPhysicalPrimary(mclabel))
        mctype += 1;
      if(fMCStack->IsSecondaryFromWeakDecay(mclabel))
        mctype += 2;
      if(fMCStack->IsSecondaryFromMaterial(mclabel))
        mctype += 4;
    }
    
    Int_t run = fEvent->GetRunNumber();

    (*fStream)<<"tree"<<
      "ifilled="<<ifilled<<
      "run="<<run<<
      //"cutbit="<<cutbit[itrk]<<

      "v0cen="<<v0cen<<

      //"dcaxy="<<dcaxy<<
      //"dcaz="<<dcaz<<
      //"tpcchi2cg="<<tpcchi2cg<<

      "mom="<<mom<<
      "charge="<<charge<<
      "pt="<<pt<<
      "eta="<<eta<<

      "tofpion="<<tofpion<<
      "tofkaon="<<tofkaon<<
      "tofproton="<<tofproton<<
      "tofelectron="<<tofelectron<<
      //"tofmass="<<tofmass<<

      "tpcsig="<<tpcsig<<
      //"tpcpidn="<<tpcpidn<<
      //"tpcncls="<<tpcncls<<
      //"tpcns="<<tpcns<<
      //"tpcnf="<<tpcnf<<
      //"tpcncross="<<tpcncross<<
      "tpcip="<<tpcip;

    if(fMCStack){
      (*fStream)<<"tree"<<
        "mcpdg="<<mcpdg<<
        "mctype="<<mctype<<
        "mcpt="<<mcpt<<
        "mceta="<<mceta<<
        "mccharge="<<mccharge;
    }

    (*fStream)<<"tree"<<
      "\n";
    
    ifilled++;
  } 
  if(ifilled){
    fCuts->Fill(5., v0cen); //================================
    fNevent->Fill(1);
  }
}

void AliAnalysisTaskpkRaa::SaveMCgen(const Float_t tmpcen, const Int_t tmpkev)
{
  //http://www.slac.stanford.edu/BFROOT/www/Computing/Environment/NewUser/htmlbug/node51.html
  const Int_t kkaon = 321;

  Float_t v0cen = tmpcen;
  Int_t kFailEvent = tmpkev;

  for(Int_t ii=0; ii<fMCStack->GetNtrack(); ii++){
    TParticle * trackMC = fMCStack->Particle(ii);
    Int_t mcpdg = TMath::Abs(trackMC->GetPdgCode());

    //remove other particel types because their amount depends on the physics in MC (namely the generator), not always corresponding to real data              
    if(mcpdg!=kElectron && mcpdg!=kProton && mcpdg != kPiPlus && mcpdg != kkaon)
      continue;

    //this is pseudo-rapidity
    //http://root.cern.ch/root/html/src/TParticle.h.html#AlKnRB
    Float_t mceta = trackMC->Eta();
    if(!EtaCut(mceta))
      continue;
    
    Int_t mccharge = -999;
    //root [11] TParticlePDG::Charge(
    //Double_t Charge() const         // charge in units of |e|/3
    if(trackMC->GetPDG()->Charge()>0)
      mccharge = 1;
    else if(trackMC->GetPDG()->Charge()<0)
      mccharge = -1;
    else 
      mccharge = 0;

    Float_t mcpt = trackMC->Pt();

    Int_t mctype = 0;
    if(fMCStack->IsPhysicalPrimary(ii))
      mctype += 1;
    if(fMCStack->IsSecondaryFromWeakDecay(ii))
      mctype += 2;
    if(fMCStack->IsSecondaryFromMaterial(ii))
      mctype += 4;

    (*fMCStream)<<"mctree"<<
      "kFailEvent="<<kFailEvent<<
      "v0cen="<<v0cen<<
      "mcpdg="<<mcpdg<<
      "mceta="<<mceta<<
      "mccharge="<<mccharge<<
      "mcpt="<<mcpt<<
      "mctype="<<mctype<<
      "\n";
  }
}

Bool_t AliAnalysisTaskpkRaa::EtaCut(const Float_t eta)
{
  if(TMath::Abs(eta)<0.8)
    return kTRUE;

  return kFALSE;
}

TString AliAnalysisTaskpkRaa::GetRunInfo()
{
  const Int_t run = fEvent->GetRunNumber();
  //checked on E-log ->
  //10a:        -> 114649
  //10b: 114650 -> 117630
  //10c: 117631 -> 121526
  //10d: 121527 -> 126460
  //10e: 126461 -> 130930
  //10f: 130931 ->

  //10g:        -> 136781
  //10h: 136782 -> 139846
  //11a: 139847 -> 146974
  //11b: 146975

  //11f:        -> 165771 
  //11h: 165772 -> 170718
  //12a: 170719 ->
  //<-

  if(     run >= 114650 && run <= 117630 ){
    return "pp10b";
  }
  else if(run >= 117631 && run <= 121526 ){
    return "pp10c";
  }
  else if(run >= 121527 && run <= 126460 ){
    return "pp10d";
  }
  else if(run >= 126461 && run <= 130930){
    return "pp10e";
  }
  else if(run >= 136782 && run <= 139846){
    return "PbPb10h";
  }
  else if(run >= 139847 && run <= 146974){
    return "pp11a";
  }
  else if(run >= 165772 && run <= 170718){
    return "PbPb11h";
  }
  else{
    printf("AliAnalysisTaskpkRaa::GetRunInfo wrong run! %d\n", run);exit(1);
  }
}

Bool_t AliAnalysisTaskpkRaa::CutInput()
{
  //
  //general protection
  //

  //MC

  fMCEvent = 0x0;
  fMCStack = 0x0;

  AliMCEventHandler* mcH = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
  if(mcH){
    fMCEvent = mcH->MCEvent();
    if(!fMCEvent){
      printf("AliAnalysisTaskpkRaa::CutInput no mcevent from non-0 mchandler!\n");exit(1);
    }
    fMCStack = fMCEvent->Stack();
    if(!fMCStack){
      printf("AliAnalysisTaskpkRaa::CutInput no mcstack from non-0 mcevent!\n");exit(1);
    }
  }

  //================================================
  //ESD

  fEvent = 0x0;
  fpid = 0x0;
  AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler *>(AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  if(esdH){  
    fEvent = esdH->GetEvent();
    fpid = esdH->GetPIDResponse();
  }

  if(!fEvent){
    printf("AliAnalysisTaskpkRaa::CutInput error No ESD Event\n"); exit(1);
  }
  
  if(!fpid){
    printf("AliAnalysisTaskpkRaa::CutInput error No ESD ESDpid\n"); exit(1);
  }

  //in unit of kg
  const Float_t dm = fEvent->GetMagneticField();
  if( TMath::Abs(TMath::Abs(dm)- 5) > 0.1 ){
    printf("AliAnalysisTaskpkRaa::CutInput error strange Bfield! %f\n", dm); exit(1);
  }

  return kTRUE;
}

Int_t AliAnalysisTaskpkRaa::FailCutEvent(const Float_t v0cen)
{
  fCuts->Fill(10., v0cen);

  const AliESDVertex *primVertex = fEvent->GetPrimaryVertexTracks();
  if(!primVertex){
    return 1;
  }

  fCuts->Fill(11., v0cen);

  if(primVertex->GetNContributors()<1) {
    // SPD vertex
    primVertex = fEvent->GetPrimaryVertexSPD();
    if(primVertex->GetNContributors()<1) {
      // NO GOOD VERTEX, SKIP EVENT 
      return 2;
    }
  }

  fCuts->Fill(12., v0cen);

  fVtxZ->Fill(primVertex->GetZ());
  primVertex = fEvent->GetPrimaryVertex();
  if(TMath::Abs(primVertex->GetZ())>10 ){
    return 3;
  }

  fCuts->Fill(13., v0cen);

  return 0;
}

void AliAnalysisTaskpkRaa::CutTrack(Int_t &nsel,  typetrack *outtrk[], Int_t  cutbit[])
{
  //
  //CutTrack
  //

  //only use RAA final cut so that the number of event is correct!
  //==========================================================================================
  //==========================================================================================
  //keep it to check whether opencut + addtional cut can reproduce raa cut

  //https://alimonitor.cern.ch/users/download.jsp?view=true&path=/alice/cern.ch/user/a/alidaq/AOD/AOD086/AddTaskESDFilter.C
  // standard cuts also used in R_AA analysis
  //   "Global track RAA analysis QM2011 + Chi2ITS<36";
  // 1024 1<<10        
  //also used for 11a pp data
  AliESDtrackCuts *raaCut = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE,1);
  raaCut->SetMinNCrossedRowsTPC(120);
  raaCut->SetMinRatioCrossedRowsOverFindableClustersTPC(0.8);
  raaCut->SetMaxChi2PerClusterITS(36);
  raaCut->SetMaxFractionSharedTPCClusters(0.4);
  raaCut->SetMaxChi2TPCConstrainedGlobal(36);
  raaCut->SetEtaRange(-0.9,0.9);
  raaCut->SetPtRange(0.15, 1e10);
  
  //==========================================================================================
  //==========================================================================================

  nsel = 0;
  for(Int_t ii=0; ii<fEvent->GetNumberOfTracks(); ii++){
    typetrack *trackptr = fEvent->GetTrack(ii);

    cutbit[nsel] = 0;

    if(raaCut->AcceptTrack(trackptr) && trackptr->GetTPCsignalN()>60){
      cutbit[nsel] += 1;
    }
   
    if(cutbit[nsel]){
      outtrk[nsel]=trackptr;
      nsel++;
    }
  }
  
  delete raaCut;
}

